package org.springframework.session.web.socket.config.annotation;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.annotation.Bean;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.session.Session;
import org.springframework.session.SessionRepository;
import org.springframework.session.web.socket.handler.WebSocketConnectHandlerDecoratorFactory;
import org.springframework.session.web.socket.handler.WebSocketRegistryListener;
import org.springframework.session.web.socket.server.SessionRepositoryMessageInterceptor;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.StompWebSocketEndpointRegistration;
import org.springframework.web.socket.config.annotation.WebMvcStompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration;
import org.springframework.web.socket.messaging.StompSubProtocolErrorHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.util.UrlPathHelper;

public abstract class AbstractSessionWebSocketMessageBrokerConfigurer<S extends Session> implements WebSocketMessageBrokerConfigurer {

	@Autowired
	@SuppressWarnings("rawtypes")
	private SessionRepository sessionRepository;

	@Autowired
	private ApplicationEventPublisher eventPublisher;

	@Override
	public void configureClientInboundChannel(ChannelRegistration registration) {
		registration.interceptors(sessionRepositoryInterceptor());
	}

	@Override
	public final void registerStompEndpoints(StompEndpointRegistry registry) {
		if (registry instanceof WebMvcStompEndpointRegistry) {
			WebMvcStompEndpointRegistry mvcRegistry = (WebMvcStompEndpointRegistry) registry;
			configureStompEndpoints(new SessionStompEndpointRegistry(mvcRegistry, sessionRepositoryInterceptor()));
		}
	}

	protected abstract void configureStompEndpoints(StompEndpointRegistry registry);

	@Override
	public void configureWebSocketTransport(WebSocketTransportRegistration registration) {
		registration.addDecoratorFactory(wsConnectHandlerDecoratorFactory());
	}

	@Bean
	public WebSocketRegistryListener webSocketRegistryListener() {
		return new WebSocketRegistryListener();
	}

	@Bean
	public WebSocketConnectHandlerDecoratorFactory wsConnectHandlerDecoratorFactory() {
		return new WebSocketConnectHandlerDecoratorFactory(this.eventPublisher);
	}

	@Bean
	@SuppressWarnings("unchecked")
	public SessionRepositoryMessageInterceptor<S> sessionRepositoryInterceptor() {
		return new SessionRepositoryMessageInterceptor<>(this.sessionRepository);
	}

	static class SessionStompEndpointRegistry implements StompEndpointRegistry {

		private final WebMvcStompEndpointRegistry registry;

		private final HandshakeInterceptor interceptor;

		SessionStompEndpointRegistry(WebMvcStompEndpointRegistry registry, HandshakeInterceptor interceptor) {
			this.registry = registry;
			this.interceptor = interceptor;
		}

		@Override
		public StompWebSocketEndpointRegistration addEndpoint(String... paths) {
			StompWebSocketEndpointRegistration endpoints = this.registry.addEndpoint(paths);
			endpoints.addInterceptors(this.interceptor);
			return endpoints;
		}

		@Override
		public void setOrder(int order) {
			this.registry.setOrder(order);
		}

		@Override
		public void setUrlPathHelper(UrlPathHelper urlPathHelper) {
			this.registry.setUrlPathHelper(urlPathHelper);
		}

		@Override
		public WebMvcStompEndpointRegistry setErrorHandler(StompSubProtocolErrorHandler errorHandler) {
			return this.registry.setErrorHandler(errorHandler);
		}
	}
}
